import os
import cv2
import mycv
import utils
import random
import RotateImg
from tqdm import tqdm


path_src = r'/media/fang/TOSHIBA EXT/Temp/Big_M1_save2'
path_output = r'/media/fang/TOSHIBA EXT/Temp/Big_Balance'

labels = ['0', '10', '11', '12', '13', '15', '2', '3', '4', '5', '6', '7', '8', '9', 'Limit', 'RemoveLimit']
num_src = [10661, 910, 708, 481, 624, 647, 1540, 1637, 1252, 3958, 1518, 1497, 1696, 286, 4736, 7832]
num_need = 20000
num_src[-1] = num_need
num_src[-2] = num_need
names = [[] for i in range(len(num_src) - 2)]
for root, dirs, files in os.walk(path_src):
    for file in tqdm(files):
        if file.split(".")[-1] == "xml":
            infos = utils.ReadXml(root + '/' + file)
            for info in infos:
                if info[-1] not in ['Limit', 'RemoveLimit']:
                    names[labels.index(info[-1])].append(root + '/' + file)

name_num = 1
All_is_ok = False
while All_is_ok is False:
    un_ok_cnt = 0
    for idx_label in range(len(labels) - 2):
        if num_src[idx_label] < num_need:
            un_ok_cnt += 1
            file = random.choice(names[idx_label])
            infos = utils.ReadXml(file)
            for info in infos:
                num_src[labels.index(info[-1])] += 1
            img = cv2.imread(file.replace('xml', 'jpg'))
            img, infos = RotateImg.RotateImgXml(img, infos, [-5, 5])
            img = mycv.motion_blur(img, random.randint(1, 5))
            infos.insert(0, path_output + '/B%08d.xml' % name_num)
            name_num += 1
            utils.WriteXml(infos, img.shape[1], img.shape[0])
            cv2.imwrite(infos[0].replace('xml', 'jpg'), img)
    if un_ok_cnt == 0:
        All_is_ok = True
